import pandas as pd, numpy as np
from pathlib import Path
from sklearn.model_selection import train_test_split
from sklearn.metrics import roc_auc_score, precision_recall_fscore_support, confusion_matrix
from xgboost import XGBClassifier

# 1. 读数据（表二）
path = Path("data_train_balanced.xlsx")
xls = pd.ExcelFile(path, engine='openpyxl')
df = pd.read_excel(xls,  engine='openpyxl')



# 3. 特征
features = [
    '13号染色体的Z值','18号染色体的Z值','21号染色体的Z值','X染色体的Z值',
    'X染色体浓度','13号染色体的GC含量','18号染色体的GC含量','21号染色体的GC含量',
    '孕妇BMI','检测孕周','原始读段数','在参考基因组上比对的比例',
    '重复读段的比例','被过滤掉读段数的比例'
]
features = [f for f in features if f in df.columns]

# 转换孕周为数值
def parse_week(s):
    try:
        s = str(s)
        if 'w' in s:
            parts = s.split('w')
            weeks = int(parts[0])
            days = 0
            if '+' in parts[1]:
                days = int(parts[1].replace('+',''))
            return weeks + days/7.0
        else:
            return float(s)
    except:
        return np.nan
if '检测孕周' in df.columns:
    df['检测孕周'] = df['检测孕周'].apply(parse_week)

X = df[features].apply(pd.to_numeric, errors='coerce')
X = X.fillna(X.median())
y = df['是否异常'].astype(int)

# 4. 按孕妇代码分训练/测试
patients = df['孕妇代码'].astype(str) if '孕妇代码' in df.columns else df.index.astype(str)
patient_label = pd.DataFrame({'patient': patients, 'label': y})
patient_label = patient_label.groupby('patient')['label'].max().reset_index()
train_patients, test_patients = train_test_split(
    patient_label['patient'], test_size=0.2, stratify=patient_label['label'], random_state=42
)
train_mask = patients.isin(train_patients)
test_mask = patients.isin(test_patients)
X_train, X_test = X[train_mask], X[test_mask]
y_train, y_test = y[train_mask], y[test_mask]

# 5. 计算 scale_pos_weight 处理不平衡 - 为了提高recall，增加正样本权重
scale_pos_weight = ((y_train == 0).sum() / (y_train == 1).sum()) * 2.0  # 增加2倍权重

# 6. 训练 XGBoost 分类器 - 优化参数以提高recall
xgb = XGBClassifier(
    n_estimators=200,        # 增加树的数量
    learning_rate=0.05,      # 降低学习率以提高稳定性
    max_depth=8,             # 适中的树深度
    min_child_weight=1,      # 降低最小子节点权重
    subsample=0.8,
    colsample_bytree=0.8,
    scale_pos_weight=scale_pos_weight,
    objective="binary:logistic",
    random_state=42,
    n_jobs=-1,
    eval_metric="logloss"    # 使用logloss更适合概率预测
)
xgb.fit(X_train, y_train)

# 7. 阈值优化分析
probs = xgb.predict_proba(X_test)[:,1]

print("=== 阈值优化分析 ===")
print("阈值\tPrecision\tRecall\t\tF1\t\tTP\tFP\tTN\tFN")
print("-" * 70)

thresholds = [0.1, 0.15, 0.2, 0.25, 0.3, 0.35, 0.4, 0.45, 0.5, 0.55, 0.6]
best_recall = 0
best_threshold = 0.3
best_f1 = 0
optimal_threshold = 0.3

for threshold in thresholds:
    preds = (probs >= threshold).astype(int)
    prec, rec, f1, _ = precision_recall_fscore_support(y_test, preds, average='binary', zero_division=0)
    cm = confusion_matrix(y_test, preds)
    tn, fp, fn, tp = cm.ravel()
    
    print(f"{threshold:.2f}\t{prec:.4f}\t\t{rec:.4f}\t\t{f1:.4f}\t\t{tp}\t{fp}\t{tn}\t{fn}")
    
    # 优先选择recall高的阈值，在recall相同时选择F1更高的
    if rec > best_recall or (rec == best_recall and f1 > best_f1):
        best_recall = rec
        best_f1 = f1
        best_threshold = threshold
        optimal_threshold = threshold

print(f"\n最优阈值: {optimal_threshold}")
print(f"该阈值下的Recall: {best_recall:.4f}")
print(f"该阈值下的F1: {best_f1:.4f}")

# 使用最优阈值进行最终预测
final_preds = (probs >= optimal_threshold).astype(int)
auc = roc_auc_score(y_test, probs)
final_prec, final_rec, final_f1, _ = precision_recall_fscore_support(y_test, final_preds, average='binary', zero_division=0)
final_cm = confusion_matrix(y_test, final_preds)

print(f"\n=== 使用最优阈值 {optimal_threshold} 的最终结果 ===")
print("AUC:", auc)
print("Precision:", final_prec)
print("Recall:", final_rec)
print("F1:", final_f1)
print("Confusion Matrix:\n", final_cm)


new_data_path = Path("data_test_validation.xlsx")
xls_new = pd.ExcelFile(new_data_path, engine="openpyxl")
df_new = pd.read_excel(xls_new, engine="openpyxl")



# 特征处理
if '检测孕周' in df_new.columns:
    df_new['检测孕周'] = df_new['检测孕周'].apply(parse_week)
X_new = df_new[features].apply(pd.to_numeric, errors="coerce")
X_new = X_new.fillna(X.median())   # 注意：用旧数据的中位数填补
y_new = df_new['是否异常'].astype(int)

# 使用最优阈值进行新数据预测
probs_new = xgb.predict_proba(X_new)[:,1]
preds_new = (probs_new >= optimal_threshold).astype(int)

# 评估
auc_new = roc_auc_score(y_new, probs_new)
prec_new, rec_new, f1_new, _ = precision_recall_fscore_support(y_new, preds_new, average="binary", zero_division=0)
cm_new = confusion_matrix(y_new, preds_new)

print(f"\n=== 在新数据上使用最优阈值 {optimal_threshold} 的表现 ===")
print("AUC:", auc_new)
print("Precision:", prec_new)
print("Recall:", rec_new)
print("F1:", f1_new)
print("混淆矩阵:\n", cm_new)